import torch
from PIL import Image
from matplotlib import pyplot as plt
from torch.utils.data import DataLoader
import numpy as np
from lib.utils.dataloader import FacenetDataset, dataset_collate
from lib.utils.utils import resize_image, cvtColor, preprocess_input

input_shape = [160, 160, 3]
with open(r'../dataset/train.txt', "r") as f:
    lines = f.readlines()
#
# train_dataset = FacenetDataset(input_shape, lines, random=True)
#
# # 用以训练的数据
# gen = DataLoader(train_dataset, shuffle=True, batch_size=96 // 2, num_workers=4,
#                  pin_memory=True, drop_last=True, collate_fn=dataset_collate, sampler=None)
#
# if __name__ == '__main__':
#     for iteration, batch in enumerate(gen):
#         images, labels = batch
#         print("images.shape=", images.shape)
#         print("labels.shape=", labels.shape)
# imageA = cvtColor(Image.open(r'../dataset/train/data/377/a.jpg'))
# print(imageA.size)
#
# imageB = resize_image(imageA, [input_shape[0], input_shape[1]], letterbox_image=True)
# print(imageB.size)
#
# plt.subplot(1, 2, 1)
# plt.imshow(np.array(imageA))
#
# plt.subplot(1, 2, 2)
# plt.imshow(np.array(imageB))
#
# plt.show()


a = FacenetDataset(input_shape=input_shape, imginfo=lines, random=True)
d = a.dataDic
print(d['377'])




